import numpy as np


class TradeAccount:
    """管理交易账户的状态和操作"""

    def __init__(
            self,
            initial_balance=1000000,         # 初始账户资金
            leverage=100,                    # 杠杆倍数，用于计算保证金
            fee_rate=0.0005,                  # 交易手续费率
            max_position_ratio= 0.00001,  # 最大持仓比例
    ):
        # 账户参数
        self.initial_balance = initial_balance
        self.leverage = leverage
        self.fee_rate = fee_rate
        self.max_position_ratio = max_position_ratio

        # 账户状态
        self.balance = initial_balance  # 当前余额
        self.net_worth = initial_balance  # 当前净值
        self.old_net_worth = initial_balance
        self.max_net_worth = initial_balance  # 历史最大净值
        self.long_position = 0.0  # 多头持仓量
        self.long_ave_price = 0.0  # 多头平均持仓价
        self.short_position = 0.0  # 空头持仓量
        self.short_ave_price = 0.0  # 空头平均持仓价
        self.current_price = None

        self.last_long_step = 0 # 上一次多头操作的步数
        self.last_short_step = 0 # 上一次空头操作的步数
        self.order_nums = 0 # 当前下单次数
        self.loss_order_nums = 0 # 当前亏损下单次数
        self.win_order_nums = 0 # 当前盈利下单次数
        self.today = 0 # 当前日期
        self.steps = 0 # 当前步数
    def reset(self):
        """重置账户状态到初始值"""
        self.balance = self.initial_balance
        self.net_worth = self.initial_balance
        self.old_net_worth = self.initial_balance
        self.max_net_worth = self.initial_balance
        self.long_position = 0.0
        self.long_ave_price = 0.0
        self.short_position = 0.0
        self.short_ave_price = 0.0
        self.loss_order_nums = 0
        self.win_order_nums = 0
        self.current_price = None

        self.last_long_step = 0
        self.last_short_step = 0
        self.order_nums = 0
        self.today = 0
        self.steps = 0

    def set_price(self, current_price):
        self.current_price =current_price
        self.steps += 1
        self.today = (self.steps // 86_400) + 1
    def get_max_amount(self):
        return self.initial_balance / self.current_price * self.max_position_ratio * self.leverage
    def get_fee(self, amount):
        return amount * self.current_price * self.fee_rate
    def get_pay_balance(self, amount):
        return amount * self.current_price / self.leverage + self.get_fee(amount)
    def open_long(self):
        if not self.is_can_open_long():
            return False

        current_price = self.current_price
        amount = self.get_max_amount()
        cost = self.get_pay_balance(amount)

        if self.balance >= cost and amount > 0:
            total_cost = self.long_position * self.long_ave_price + amount * current_price
            self.long_position += amount
            self.long_ave_price = total_cost / self.long_position
            self.balance -= cost
            self.order_nums += 1
            self.last_long_step = self.steps
            return -self.get_fee(amount)
        return False

    def close_long(self):
        if not self.is_can_close_long():
            return False
        current_price = self.current_price

        amount = self.long_position
        fee = self.get_fee(amount)
        profit = (current_price - self.long_ave_price) * amount - fee
        margin = self.long_ave_price * amount / self.leverage

        self.balance += margin + profit
        self.long_position -= amount
        if self.long_position <= 0:
            self.long_ave_price = 0.0
        if profit <= 0:
            self.loss_order_nums += 1
        else:
            self.win_order_nums += 1
        return profit

    def open_short(self):
        if not self.is_can_open_short():
            return False
        current_price = self.current_price
        amount = self.get_max_amount()
        cost = self.get_pay_balance(amount)

        if self.balance >= cost and amount > 0:
            total_cost = self.short_position * self.short_ave_price + amount * current_price
            self.short_position += amount
            self.short_ave_price = total_cost / self.short_position
            self.balance -= cost
            self.order_nums += 1
            self.last_short_step = self.steps
            return -self.get_fee(amount)
        return False
    def close_short(self):
        if not self.is_can_close_short():
            return False
        current_price = self.current_price

        amount = self.short_position
        fee = self.get_fee(amount)
        profit = (self.short_ave_price - current_price) * amount - fee
        margin = self.short_ave_price * amount / self.leverage

        self.balance += margin + profit
        self.short_position -= amount
        if self.short_position <= 0:
            self.short_ave_price = 0.0
        if profit <= 0:
            self.loss_order_nums += 1
        else:
            self.win_order_nums += 1
        return profit

    def close_all(self):
        self.close_short()
        self.close_long()

    def update_net_worth(self):
        current_price = self.current_price

        long_margin = self.long_position * self.long_ave_price / self.leverage
        short_margin = self.short_position * self.short_ave_price / self.leverage

        long_pnl = self.long_position * (current_price - self.long_ave_price)
        short_pnl = self.short_position * (self.short_ave_price - current_price)

        self.old_net_worth = self.net_worth
        self.net_worth = self.balance + long_margin + short_margin + long_pnl + short_pnl
        self.max_net_worth = max(self.max_net_worth, self.net_worth)

    def get_account_state(self):
        # 基础账户状态：余额、多空持仓、平均持仓价
        base_state = np.array([
            self.balance,
            self.net_worth,
            self.max_net_worth,
            self.long_position,
            self.short_position,
            self.long_ave_price,
            self.short_ave_price,
            self.get_drawdown(),
            self.get_gain_ratio(),
        ], dtype=np.float32)

        return np.concatenate([base_state])
    def is_can_open_long(self):
        return self.long_position <= 0 and self.short_position <= 0
    def is_can_close_long(self):
        return self.long_position > 0
    def is_can_open_short(self):
        return self.short_position <= 0 and self.long_position <= 0
    def is_can_close_short(self):
        return self.short_position > 0

    def is_holding(self):
        """检查是否有持仓"""
        return self.long_position > 0 or self.short_position > 0
    def is_holding_long(self):
        """检查是否有多头持仓"""
        return self.long_position > 0
    def is_holding_short(self):
        """检查是否有空头持仓"""
        return self.short_position > 0
    def get_long_steps(self):
        if self.long_position <= 0:
            return 0
        return self.steps - self.last_long_step
    def get_short_steps(self):
        if self.short_position <= 0:
            return 0
        return self.steps - self.last_short_step

    def get_drawdown(self):
        """计算当前回撤"""
        return (self.net_worth - self.max_net_worth) / self.max_net_worth

    def get_gain_ratio(self):
        """计算账户盈亏比例"""
        return (self.net_worth - self.initial_balance) / self.initial_balance

    def remain_steps(self): # 计算剩余步数比例
        return (86_400 - self.steps % 86_400)/86_400
